import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Dropout
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
import os

def create_lstm_model(sequence_length, n_features=1, units=50):
    """
    创建LSTM模型
    
    参数:
        sequence_length: 输入序列长度
        n_features: 特征数量
        units: LSTM单元数量
        
    返回:
        编译好的LSTM模型
    """
    model = Sequential([
        LSTM(units, activation='relu', input_shape=(sequence_length, n_features), return_sequences=False),
        Dropout(0.2),
        Dense(1)
    ])
    
    model.compile(optimizer='adam', loss='mse')
    return model

def train_lstm_model(model, train_generator, test_generator, epochs=50, model_path=None):
    """
    训练LSTM模型
    
    参数:
        model: 要训练的模型
        train_generator: 训练数据生成器
        test_generator: 测试数据生成器
        epochs: 训练轮数
        model_path: 模型保存路径
        
    返回:
        训练历史和训练好的模型
    """
    callbacks = [
        EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True)
    ]
    
    if model_path:
        os.makedirs(os.path.dirname(model_path), exist_ok=True)
        callbacks.append(ModelCheckpoint(model_path, save_best_only=True))
    
    history = model.fit(
        train_generator,
        epochs=epochs,
        validation_data=test_generator,
        callbacks=callbacks
    )
    
    return history, model

def predict_lstm(model, test_generator, test_data, sequence_length):
    """
    使用LSTM模型进行预测
    
    参数:
        model: 训练好的模型
        test_generator: 测试数据生成器
        test_data: 测试数据
        sequence_length: 序列长度
        
    返回:
        预测结果
    """
    predictions = []
    
    # 使用生成器进行预测
    batch_predictions = model.predict(test_generator)
    predictions.extend(batch_predictions.flatten())
    
    # 将预测结果转换为适当的形状
    predictions = np.array(predictions).reshape(-1, 1)
    
    # 创建与测试数据相同长度的预测数组，前sequence_length个值为NaN
    full_predictions = np.full(test_data.shape, np.nan)
    full_predictions[sequence_length:sequence_length+len(predictions)] = predictions
    
    return full_predictions